# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import getpass
import logging
import math
import os
import subprocess
import sys
import tempfile

from absl import app, flags

from coin.base.datetime_util import to_datetime
from coin.base.param_util import to_list
from coin.strategy.mm.simple_sim.runner_util import (get_profile_module, get_num_strategies)


def datetime_to_str(dt):
  dt = to_datetime(dt)
  return dt.strftime('%Y%m%d-%H%M%S')


class NoJobError(Exception):
  pass


def generate_slurm_script(profile_name,
                          profile_args,
                          max_strategy_per_process=None,
                          job_submit_period=None,
                          outdir_prefix=None,
                          user=None,
                          grid=None,
                          venv=None,
                          summary_mailto=None,
                          flow_root_dir=None,
                          tag=None):
  max_strategy_per_process = max_strategy_per_process or 16
  job_submit_period = 0.01 if job_submit_period is None else job_submit_period
  user = user or getpass.getuser()
  grid = grid or 'iosg'
  outdir_prefix = (outdir_prefix
                   or '/remote/{grid}/home/{user}/scratch'.format(grid=grid, user=user))

  profile = get_profile_module(profile_name, profile_args)
  time_ranges = to_list(profile.get_time_ranges())

  lines = ['#!/bin/bash', '#SBATCH --spread-job', '# Copyright (c) 2018 Presto Labs Pte. Ltd.\n']

  lines += [
      'export OUTDIR="%s/sim-${SLURM_JOB_ID}"' % outdir_prefix,
      'mkdir -p "${OUTDIR}"\n',
      'if [ -z ${SLURM_JOB_ID} ]; then',
      '  (>&2 echo "ERROR: run this script on slurm.")',
      '  exit 1',
      'fi\n',
      'set -e'
  ]

  if venv is not None:
    lines.append('source "%s/bin/activate"' % venv)

  if tag:
    lines.append('printf "Tag: %s\\n\\n" > "${OUTDIR}/summary.txt"' % tag)

  lines.append('\nNUM_PARALLEL="$((${SLURM_NTASKS} + 30))"')
  lines.append('parallel -j "${NUM_PARALLEL}" --delay %f <<EOF' % job_submit_period)

  out_idx = 0
  for from_ts, to_ts in time_ranges:
    num_strat = get_num_strategies(profile, from_ts, to_ts)
    num_group = int(math.ceil(num_strat / float(max_strategy_per_process)))
    if num_group == 0:
      continue
    step = int(math.ceil(num_strat / float(num_group)))

    for i in range(num_group):
      from_idx = i * step
      to_idx = min(from_idx + step - 1, num_strat - 1)
      if to_idx < from_idx:
        continue

      cmd = [
          'srun',
          '-n1',
          '-l',
          '--exclusive',
          'python',
          '-m',
          'coin.strategy.mm.simple_sim.slurm.single_process_runner',
          '--profile="%s"' % profile_name,
          '--timerange_from_ts="%s"' % datetime_to_str(from_ts),
          '--timerange_to_ts="%s"' % datetime_to_str(to_ts),
          '--strategy_from_idx=%06d' % from_idx,
          '--strategy_to_idx=%06d' % to_idx,
          '--out="${OUTDIR}/result-%06d.json"' % out_idx
      ]
      if flow_root_dir:
        cmd += ['--flow_root_dir="%s"' % flow_root_dir]
      cmd += ['--', *profile_args]
      out_idx += 1

      lines.append(' '.join(cmd))

  if out_idx == 0:
    raise NoJobError()

  lines.append('EOF\n')
  sys.stderr.write('# Sim processes: %d\n' % (out_idx + 1))

  result_aggregator_cmd = [
      'python',
      '-m',
      'coin.strategy.mm.simple_sim.slurm.result_aggregator',
      '--profile="%s"' % profile_name,
      '--result_dir="${OUTDIR}"',
      '--',
      *profile_args,
      '>> "${OUTDIR}/summary.txt"'
  ]
  lines.append(' '.join(result_aggregator_cmd))

  if summary_mailto is not None:
    if tag:
      mail_subject = "[coin-sim] ${SLURM_JOB_ID} summary - %s" % tag
    else:
      mail_subject = "[coin-sim] ${SLURM_JOB_ID} summary"

    mail_cmd = [
        'python',
        '-m',
        'coin.tool.mailer',
        '--name="coin-sim"',
        '--mailfrom="coin-sim@joomo.io"',
        '--mailto="%s"' % summary_mailto,
        '--subject="%s"' % mail_subject,
        '--attach="${OUTDIR}/summary.txt"',
        '<<EOF'
    ]

    lines += [
        ' '.join(mail_cmd),
        'Slurm job ID:     ${SLURM_JOB_ID}',
        'Output directory: ${OUTDIR}',
        'Summary file:     ${OUTDIR}/summary.txt',
        'User tag:         %s\n\n' % tag,
        '---',
        '$(scontrol show job ${SLURM_JOB_ID})',
        'EOF'
    ]

  return '\n'.join(lines)


class SBatchError(Exception):
  pass


def run_sbatch(script_filepath, sbatch_args=None, verbose=False):
  sbatch_args = sbatch_args or []
  if isinstance(sbatch_args, str):
    sbatch_args = [token for token in sbatch_args.split(' ') if token]

  sbatch_args = ['--parsable'] + sbatch_args
  if sbatch_args:
    logging.debug('sbatch_args: %s', ' '.join(sbatch_args))

  completed = subprocess.run(['sbatch', *sbatch_args, script_filepath], stdout=subprocess.PIPE)
  if completed.returncode != 0:
    raise SBatchError(completed)

  sbatch_output = completed.stdout.decode('utf-8')
  job_id = int(sbatch_output.strip())
  if verbose:
    sys.stderr.write('Job ID: %d\n' % job_id)
  return job_id


def run_simulator_on_slurm(profile_name,
                           profile_args,
                           max_strategy_per_process=None,
                           job_submit_period=None,
                           outdir_prefix=None,
                           user=None,
                           grid=None,
                           venv=None,
                           force_use_venv=True,
                           summary_mailto='{user}@prestolabs.io',
                           tag=None,
                           flow_root_dir=None,
                           dry_run=False,
                           sbatch_args=None,
                           verbose=False):
  if isinstance(profile_args, tuple):
    profile_args = list(profile_args)

  if force_use_venv and venv is None:
    user = user or getpass.getuser()
    grid = grid or 'iosg'
    venv = os.environ.get('COIN_SLURM_VENV',
                          '/remote/{grid}/home/{user}/venv'.format(grid=grid, user=user))

  if summary_mailto is not None:
    if '{user}' in summary_mailto:
      user = getpass.getuser()
      summary_mailto = summary_mailto.format(user=user)
    if verbose:
      sys.stderr.write('Sim result mail will be sent to %s.\n' % summary_mailto)

  try:
    script = generate_slurm_script(profile_name,
                                   profile_args,
                                   max_strategy_per_process=max_strategy_per_process,
                                   job_submit_period=job_submit_period,
                                   outdir_prefix=outdir_prefix,
                                   user=user,
                                   grid=grid,
                                   venv=venv,
                                   summary_mailto=summary_mailto,
                                   tag=tag,
                                   flow_root_dir=flow_root_dir)
  except NoJobError:
    if verbose:
      sys.stderr.write('No job to run.\n')
    raise

  if dry_run:
    if verbose:
      print(script)
    return

  with tempfile.NamedTemporaryFile(mode='w') as script_file:
    script_file.write(script)
    script_file.flush()

    try:
      job_id = run_sbatch(script_file.name, sbatch_args, verbose=verbose)
      if verbose:
        print(job_id)
    except SBatchError:
      if verbose:
        sys.stderr.write('Error while running sbatch. Abort.\n')
      raise

  return job_id


def main(argv):
  FLAGS = flags.FLAGS

  try:
    profile_name = argv[1]
    profile_args = argv[2:]
    logging.debug('profile_args: %s', ' '.join(profile_args))
  except IndexError:
    sys.stderr.write('Usage: python -m %s.runner_slurm <profile> <profile args> -- <sbatch args>\n'
                     % __package__)
    return 1

  venv = FLAGS.venv or os.environ.get('COIN_SLURM_VENV', None)
  try:
    run_simulator_on_slurm(profile_name,
                           profile_args,
                           max_strategy_per_process=FLAGS.max_strategy_per_process,
                           outdir_prefix=FLAGS.outdir,
                           grid=FLAGS.grid,
                           venv=venv,
                           force_use_venv=False,
                           summary_mailto=FLAGS.mailto,
                           tag=FLAGS.tag,
                           flow_root_dir=FLAGS.flow_root_dir,
                           dry_run=FLAGS.script_only,
                           sbatch_args=FLAGS.sbatch_args,
                           verbose=True)
  except NoJobError:
    return 0
  except SBatchError:
    return 1


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)

  flags.DEFINE_integer('max_strategy_per_process',
                       16,
                       'Maximum number of strategies in one process')
  flags.DEFINE_string('outdir',
                      None,
                      'Output directory. Default is /remote/<grid>/home/<user>/scratch')
  flags.DEFINE_string('grid', 'iosg', 'Specify slurm grid you will use.')
  flags.DEFINE_string(
      'venv',
      None,
      'If specified, use the given virtual env. '
      'User can use venv by setting COIN_SLURM_VENV environment variable.')
  flags.DEFINE_string('mailto',
                      '{user}@prestolabs.io',
                      'If specified, it will send a mail after sim finishes.')
  flags.DEFINE_string('tag', None, 'If specified, tag will be in a result mail subject')
  flags.DEFINE_bool('script_only',
                    False,
                    'If true, it will print sbatch script and not run sbatch.')
  flags.DEFINE_string('sbatch_args', '-n 20', 'sbatch arguments')
  flags.DEFINE_string('flow_root_dir',
                      None,
                      'If specified, use the specified directory as flow roor directory.')
  app.run(main)
